package com.dataexa.ky.rest.mxgraph;

import com.dataexa.ky.constants.Constants;
import com.dataexa.ky.core.mxgraph.canvas.mxGraphicsCanvas2D;
import com.dataexa.ky.core.mxgraph.canvas.mxICanvas2D;
import com.dataexa.ky.core.mxgraph.reader.mxSaxOutputHandler;
import com.dataexa.ky.core.mxgraph.util.mxUtils;
import com.mxpdf.text.Document;
import com.mxpdf.text.DocumentException;
import com.mxpdf.text.pdf.PdfWriter;
import jakarta.ws.rs.*;
import jakarta.ws.rs.core.Context;
import jakarta.ws.rs.core.MediaType;
import jakarta.ws.rs.core.Response;
import jakarta.ws.rs.core.UriInfo;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.microprofile.openapi.annotations.Operation;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.XMLReader;

import javax.imageio.ImageIO;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.parsers.SAXParserFactory;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.StringReader;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.util.Hashtable;

@Slf4j
@Path("/mxgraph")
public class MxGraphExportResource {
    private transient SAXParserFactory parserFactory = SAXParserFactory.newInstance();
    protected transient Hashtable<String, Image> imageCache = new Hashtable<>();

    // 定义常量
    private static final int MAX_REQUEST_SIZE = 10 * 1024 * 1024; // 10MB
    private static final int MAX_AREA = 10000000; // 最大图像面积

    @Context
    UriInfo uriInfo;

    @POST
    @Consumes(MediaType.APPLICATION_FORM_URLENCODED)
    @Produces(MediaType.APPLICATION_OCTET_STREAM)
    @Operation(summary = "处理图表导出请求")
    @Path("/export")
    public Response export(
            @FormParam("format") String format,
            @FormParam("filename") String fname,
            @FormParam("w") int w,
            @FormParam("h") int h,
            @FormParam("bg") String bgParam,
            @FormParam("xml") String xmlParam) {

        try {
            // 检查请求大小
            if (xmlParam != null && xmlParam.length() > MAX_REQUEST_SIZE) {
                return Response.status(Response.Status.REQUEST_ENTITY_TOO_LARGE).build();
            }

            String xml = getRequestXml(xmlParam);

            Color bg = (bgParam != null) ? mxUtils.parseColor(bgParam) : null;

            // 检查参数
            if (w > 0 && h > 0 && w * h < MAX_AREA && format != null && xml != null && !xml.isEmpty()) {
                if (bg == null && !"png".equals(format)) {
                    bg = Color.WHITE;
                }

                if (fname != null && fname.toLowerCase().endsWith(".xml")) {
                    fname = fname.substring(0, fname.length() - 4) + "." + format;
                }

                if (fname == null) {
                    fname = "export." + format;
                }

                // 获取当前请求的URL
                String baseUrl = uriInfo.getAbsolutePath().toString();
                String url = baseUrl.substring(0, baseUrl.lastIndexOf("/"));

                ByteArrayOutputStream baos = new ByteArrayOutputStream();

                if ("pdf".equals(format)) {
                    writePdf(url, fname, w, h, bg, xml, baos);
                    return Response.ok(baos.toByteArray())
                            .header("Content-Type", "application/pdf")
                            .header("Content-Disposition", "attachment; filename=\"" + fname + "\"")
                            .build();
                } else {
                    writeImage(url, format, fname, w, h, bg, xml, baos);
                    return Response.ok(baos.toByteArray())
                            .header("Content-Type", "image/" + format.toLowerCase())
                            .header("Content-Disposition", "attachment; filename=\"" + fname + "\"")
                            .build();
                }
            } else {
                return Response.status(Response.Status.BAD_REQUEST).build();
            }
        } catch (OutOfMemoryError e) {
            log.error("图表导出出现内存溢出异常", e);
            Runtime runtime = Runtime.getRuntime();
            log.info("freeMemory: {} MB, totalMemory: {} MB, maxMemory: {} MB",
                    runtime.freeMemory() / 1024.0 / 1024,
                    runtime.totalMemory() / 1024.0 / 1024,
                    runtime.maxMemory() / 1024.0 / 1024);
            return Response.status(Response.Status.INTERNAL_SERVER_ERROR).build();
        } catch (Exception e) {
            log.error("图表导出出现未知异常", e);
            return Response.status(Response.Status.INTERNAL_SERVER_ERROR).build();
        }
    }

    protected String getRequestXml(String xml) {
        if (xml != null && xml.startsWith("%3C")) {
            xml = URLDecoder.decode(xml, StandardCharsets.UTF_8);
        }
        return xml;
    }

    protected void writeImage(String url, String format, String fname, int w, int h, Color bg, String xml, OutputStream out)
            throws IOException, SAXException, ParserConfigurationException {
        BufferedImage image = mxUtils.createBufferedImage(w, h, bg);

        if (image != null) {
            Graphics2D g2 = image.createGraphics();
            mxUtils.setAntiAlias(g2, true, true);
            renderXml(xml, createCanvas(url, g2));
            g2.dispose();
            ImageIO.write(image, format, out);
        }
    }

    protected void writePdf(String url, String fname, int w, int h, Color bg, String xml, OutputStream out)
            throws DocumentException, IOException, SAXException, ParserConfigurationException {
        // Fixes PDF offset
        w += 1;
        h += 1;

        Document document = new Document(new com.mxpdf.text.Rectangle(w, h));
        PdfWriter writer = PdfWriter.getInstance(document, out);
        document.open();

        Graphics2D graphics = writer.getDirectContent().createGraphics(w, h);
        mxGraphicsCanvas2D gc = createCanvas(url, graphics);

        // Fixes PDF offset
        gc.translate(1, 1);

        renderXml(xml, gc);
        graphics.dispose();
        document.close();
    }

    protected void renderXml(String xml, mxICanvas2D canvas) throws SAXException, ParserConfigurationException, IOException {
        XMLReader reader = parserFactory.newSAXParser().getXMLReader();
        reader.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
        reader.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);
        reader.setFeature("http://xml.org/sax/features/external-general-entities", false);
        reader.setFeature("http://xml.org/sax/features/external-parameter-entities", false);
        reader.setContentHandler(new mxSaxOutputHandler(canvas));
        reader.parse(new InputSource(new StringReader(xml)));
    }

    /**
     * Creates a graphics canvas with an image cache.
     */
    protected mxGraphicsCanvas2D createCanvas(String url, Graphics2D g2) {
        // Caches custom images for the time of the request
        final Hashtable<String, Image> shortCache = new Hashtable<String, Image>();
        final String domain = url.substring(0, url.lastIndexOf("/"));

        mxGraphicsCanvas2D g2c = new mxGraphicsCanvas2D(g2) {
            public Image loadImage(String src) {
                // Uses local image cache by default
                Hashtable<String, Image> cache = shortCache;

                // Uses global image cache for local images
                if (src.startsWith(domain)) {
                    cache = imageCache;
                }

                Image image = cache.get(src);

                if (image == null) {
                    image = super.loadImage(src);

                    if (image != null) {
                        cache.put(src, image);
                    } else {
                        cache.put(src, Constants.EMPTY_IMAGE);
                    }
                } else if (image == Constants.EMPTY_IMAGE) {
                    image = null;
                }

                return image;
            }
        };

        return g2c;
    }
}